# 抽取指定数量模型预测并渲染展示
import os
# 可视化
import random
import paddle
from PIL import Image
import numpy as np
import matplotlib.pyplot as plt
from paddleseg.core import infer
from paddleseg.datasets import Dataset
from paddleseg.models import ResNet50_vd, DeepLabV3

from src.color.color import label2color
from src.config.config import SIZE, NUM_CLASSES, MODEL_PATH, ROOT_DIR
from src.util.dataset import test_dataset, test_transforms

# 预测测试集
model = DeepLabV3(
    num_classes=NUM_CLASSES,
    backbone=ResNet50_vd(),  # currently support Resnet50_vd/Resnet101_vd/Xception65.
    pretrained=None
)

# 设置模型参数
if MODEL_PATH:
    para_state_dict = paddle.load(MODEL_PATH)
    model.set_dict(para_state_dict)
    print('Loaded trained params of model successfully')
else:
    raise ValueError('The model_path is wrong: {}'.format(MODEL_PATH))

image_dir = os.path.join(os.path.join(ROOT_DIR, '..', 'data1'))
test_dataset = Dataset(
    mode='val',  # 评估模式
    dataset_root=image_dir,
    transforms=test_transforms,
    num_classes=NUM_CLASSES,
    val_path=os.path.join(image_dir, 'val.txt'),  # 评估测试集
)

for i in [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]:
    # 随机抽取样本
    chosen_indices = random.sample(range(len(test_dataset)), k=1)
    # print(chosen_indices)

    # 读取输入影像
    image_paths = [test_dataset.file_list[idx][0] for idx in chosen_indices]

    # 获取模型预测输出
    preds = []
    with paddle.no_grad():
        model.eval()
        for idx in chosen_indices:
            image = test_dataset[idx]['img']
            # C,H,W -> 1,C,H,W
            image = paddle.to_tensor(image).unsqueeze(0)
            # 预测推理
            pred = infer.inference(model, image)
            # 对pred进行argmax操作, 并转换为NumPy数组
            pred = paddle.argmax(pred, axis=1).numpy()
            # 调整数组形状为SIZE, 并将数据类型转换为np.uint8
            pred = np.reshape(pred, SIZE).astype(np.uint8)
            # 将预测结果转为彩色
            preds.append(label2color(pred))

    # 绘制图片
    fig, axes = plt.subplots(1, 2, figsize=(1 * 5, 1 * 5))
    fig.subplots_adjust(top=1, bottom=0, left=0, right=1, hspace=0.01, wspace=0.01)

    # 给每列添加标题
    axes[0].set_title('Input Image')
    axes[0].imshow(Image.open(image_paths[0]))
    axes[0].axis("off")
    axes[0].set_title("image")
    axes[1].imshow(preds[0])
    axes[1].axis("off")
plt.show()